package workerpool

import (
	"errors"
	"fmt"
	"sync"
)

const (
	defaultCapacity int = 20
	maxCapacity     int = 50
)

var (
	ErrWorkerPoolFreed     = errors.New("worker pool has been freed")
	ErrNoAvailWorkerInPool = errors.New("no avail worker in pool")
)

type Task func()

type Pool struct {
	capacity int

	active chan struct{}
	quit   chan struct{}

	tasks chan Task

	wg sync.WaitGroup

	// 是否在创建 workerpool 时就预先创建 worker，默认值：false
	preAlloc bool

	// 当 active 已经满了，是否要阻塞 Schedule
	// 如果不阻塞，则返回 ErrNoAvaiWorkerInPool
	block bool
}

func New(capacity int, opts ...Option) *Pool {
	if capacity <= 0 {
		capacity = defaultCapacity
	}

	if capacity > maxCapacity {
		capacity = maxCapacity
	}

	p := &Pool{
		capacity: capacity,

		active: make(chan struct{}, capacity),
		quit:   make(chan struct{}),

		tasks: make(chan Task),
	}

	for _, opt := range opts {
		opt(p)
	}

	fmt.Printf("worker pool start(preAlloc=%t)\n", p.preAlloc)

	// 预先创建所有的 worker（goroutine），并发送到 channel
	if p.preAlloc {
		for i := 0; i < capacity; i++ {
			p.newWorker(i + 1)
			p.active <- struct{}{}
		}
	}

	go p.run()

	return p
}

func (p *Pool) Schedule(t Task) error {

	select {
	case <-p.quit:
		return ErrWorkerPoolFreed
	case p.tasks <- t:
		return nil
	default:
		if p.block {
			p.tasks <- t
			return nil
		}

		return ErrNoAvailWorkerInPool
	}
}

func (p *Pool) Free() {
	close(p.quit)
	p.wg.Wait()

	fmt.Println("worker pool freed")
}

func (p *Pool) run() {
	idx := len(p.active)

	if !p.preAlloc {
	loop:
		for t := range p.tasks {
			p.returnTask(t)

			select {
			case <-p.quit:
				return
			case p.active <- struct{}{}:
				idx++
				p.newWorker(idx)
			default:
				break loop
			}
		}
	}

	for {
		select {
		case <-p.quit:
			return
		case p.active <- struct{}{}:
			// create a new worker
			idx++
			p.newWorker(idx)
		}
	}
}

func (p *Pool) returnTask(t Task) {
	go func() {
		p.tasks <- t
	}()
}

func (p *Pool) newWorker(i int) {
	p.wg.Add(1)

	go func() {
		defer func() {
			// 如果发生 panic，则 active 减 1
			// 若没有发生 panic，说明是收到了 quit 信号，则 active 不减 1
			if err := recover(); err != nil {
				fmt.Printf("worker[%03d]: recover and panic exit\n", i)
				<-p.active
			}

			// 无论是否发生 panic， 都要执行 sync.WaitGroup.Done() 函数
			p.wg.Done()
		}()

		fmt.Printf("worker[%03d]: start\n", i)

		for {
			select {
			case <-p.quit:
				fmt.Printf("worker[%03d]: exit\n", i)
				<-p.active
				return
			case t := <-p.tasks:
				fmt.Printf("worker[%03d]: received a task\n", i)
				t()
			}
		}
	}()
}
